• Before proceeding refactoring, is it possible to unblock @woshiyyya? Maybe one easy solution is to build a custom wheel (you can ping reef team to learn how to do it) so that he can just use it for multi node tests.

Yes that's my plan.

• right now, we have 1 stream per recv/send for 1 process. I wonder if we need a stream per channel or should manage a queue of streams to have multiple nccl channels?

I think initially we can start with 1 stream. Whether multiple streams can provide better performance would depend on a few factors, and may need a bit of design, although we can perhaps try it.

Yeah sounds good. Pp anyway has 1:1 only. We can revisit and benchmark when we have more than 1 input output use cases

Newer commits are for graph visualization prototype, just a draft for now and no need to review new code.

Hmm there might be some misunderstanding about the suggestion to use futures. I wasn't imagining that it would need this much code change, but please let me know if my comments don't make sense.

Yeah I think the main challenge is we need to examine write values (exception check, multi-return value extraction) earlier than they are written. Now that these values are futures, we need additional handling: the problematic case is _GPUFuture that needs the send_stream to call wait() with.

Previously I came up with two approaches:

1. Transformed futures: this is the current approach with WrappedFuture, ListFuture etc.
   This approach wraps the base future and a lambda inside another future, and when the outer future is finally waited on, perform the lambda.

2. Support access to _NcclGroup._send_stream in earlier phases of the data flow (e.g., in the execution loop operations rather than nccl_group.send()), so that we can make the _send_stream wait on the future earlier
   I think this is doable if e.g., in ExecutableTask we can access the _send_stream that is to be used in send(). And I figured out a code path, the logic is roughly:

• nccl_group_id created when the nccl group is created at compile time
• Nccl_group_id is set to each TorchTensorType, which is further populated to TorchTensorNcclChannel
• nccl_group_id is used to retrieve the _NcclGroup from current_context()
• With _NcclGroup we can access _send_stream
• nccl_group_id can be passed via this path: ExecutableTask <- CompiledTask <- DAGNode <- _type_hint

I think approach2 leaves out changes to the channel interface but leaks the NcclGroup send_stream implementation details, and approach1 touches a lot of interfaces but things are consistent. Originally I thought about implementing approach2 but found approach1 may be a bit cleaner, although it changes the interfaces quite a bit.

Let me know how you think about it.

Now that these values are futures, we need additional handling: the problematic case is _GPUFuture that needs the send_stream to call wait() with.

This can be fixed by passing send_stream through the context. I saw your comment about _send_stream not being the current stream but what I meant is that we should change other parts of the code to make sure it is the current stream. I.e. right before this line, set current stream to _send_stream, _recv_stream, etc.

2. Support access to `_NcclGroup._send_stream` in earlier phases of the data flow (e.g., in the execution loop operations rather than nccl_group.send()), so that we can make the `_send_stream` wait on the future earlier

I don't understand why we need to do the wait earlier. It should be okay to just do the wait right before starting to execute the _DAGNodeOperation, same as right now. The stream wait on the event is asynchronous so it shouldn't be a performance issue.

I don't understand why we need to do the wait earlier. It should be okay to just do the wait right before starting to execute the _DAGNodeOperation, same as right now. The stream wait on the event is asynchronous so it shouldn't be a performance issue.

Yeah what I meant by "earlier" in

Support access to `_NcclGroup._send_stream` in earlier phases of the data flow (e.g., in the execution loop operations rather than nccl_group.send()), so that we can make the `_send_stream` wait on the future earlier

is exactly this line, and it needs the _send_stream to be passed via ExecutableTask, roughly with this logic I mentioned in approach2:

nccl_group_id created when the nccl group is created at compile time
nccl_group_id is set to each TorchTensorType, which is further populated to TorchTensorNcclChannel
nccl_group_id is used to retrieve the _NcclGroup from current_context()
With _NcclGroup we can access _send_stream
nccl_group_id can be passed via this path: ExecutableTask <- CompiledTask <- DAGNode <- _type_hint

So I think approach2 is what you preferred? If so, I can change to implement it.

I don't understand why we need to do the wait earlier. It should be okay to just do the wait right before starting to execute the _DAGNodeOperation, same as right now. The stream wait on the event is asynchronous so it shouldn't be a performance issue.

Yeah what I meant by "earlier" in

Support access to `_NcclGroup._send_stream` in earlier phases of the data flow (e.g., in the execution loop operations rather than nccl_group.send()), so that we can make the `_send_stream` wait on the future earlier

is exactly this line, and it needs the _send_stream to be passed via ExecutableTask, roughly with this logic I mentioned in approach2:

nccl_group_id created when the nccl group is created at compile time
nccl_group_id is set to each TorchTensorType, which is further populated to TorchTensorNcclChannel
nccl_group_id is used to retrieve the _NcclGroup from current_context()
With _NcclGroup we can access _send_stream
nccl_group_id can be passed via this path: ExecutableTask <- CompiledTask <- DAGNode <- _type_hint

So I think approach2 is what you preferred? If so, I can change to implement it.

Ah okay, got it, thanks for the explanation. Yes just to be clear, I prefer approach2 for the following reasons:

1. we don't leak details of overlapping communication etc outside of the DAGNodeOperation
2. we can call .wait() in only one place (better for testing/robustness)
3. simpler interfaces; i.e. we only need GPUFuture vs a future that's already ready, and don't need the WrappedFuture, ListFuture, etc.

lmk when ti is ready to merge! super excited for this feature.

After this can you follow up

• design for cross dag overlap
• release tests for overlap compute/comm